import os

from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr
from PIL import Image
import numpy as np
from torch import nn
from torchvision import transforms
import glob
from skimage import io

def calculateSSIM():
    print("CNN_NC")
    img3 = np.array(Image.open('CNN_NC/1true.png'))
    img4 = np.array(Image.open('CNN_NC/1pred.png'))

    sum = 0
    path = "CNN_NC"
    listName = sorted(glob.glob(os.path.join(path, '*')), key=lambda x: int((os.path.basename(x)[:-8])))

    number = len(listName) / 2
    print("一共有多少预测", number)

    result = 0
    for name in listName:
        if sum % 2 == 1:
            img4 = np.array(Image.open(name))
            result = result + ssim(img3, img4, channel_axis=2)
            sum = sum + 1
        else:
            img3 = np.array(Image.open(name))
            sum = sum + 1
    SSIM = result / number
    print("SSIM=", SSIM)

    mse = nn.MSELoss()
    img_mse1 = Image.open('CNN_NC/1true.png')
    img_mse2 = Image.open('CNN_NC/1pred.png')
    sum2 = 0
    loss = 0
    for name in listName:
        if sum2 % 2 == 1:
            img_mse1 = Image.open(name)
            ImageToTensor = transforms.ToTensor()
            tens1 = ImageToTensor(img_mse1)
            tens2 = ImageToTensor(img_mse2)
            loss = loss + mse(tens2, tens1)
            sum2 = sum2 + 1
        else:
            img_mse2 = Image.open(name)
            sum2 = sum2 + 1
    MSE = loss / number
    print("MSE=", MSE)

    mae = nn.L1Loss()
    img_mae1 = Image.open('CNN_NC/1true.png')
    img_mae2 = Image.open('CNN_NC/1pred.png')
    sum2 = 0
    loss = 0
    for name in listName:
        if sum2 % 2 == 1:
            img_mae1 = Image.open(name)
            ImageToTensor = transforms.ToTensor()
            tens1 = ImageToTensor(img_mae1)
            tens2 = ImageToTensor(img_mae2)
            loss = loss + mae(tens2, tens1)
            sum2 = sum2 + 1
        else:
            img_mae2 = Image.open(name)
            sum2 = sum2 + 1
    MAE = loss / number
    print("MAE=", MAE)

    for name in listName:
        if sum2 % 2 == 1:
            img_psnr1 = io.imread(name)

            wuhu_psnr = psnr(img_psnr1, img_psnr2)
            loss = loss + wuhu_psnr
            sum2 = sum2 + 1
        else:
            img_psnr2 = io.imread(name)
            sum2 = sum2 + 1

    PSNR = loss / number
    # 输出结果
    print("PSNR:", PSNR)
    return SSIM, MSE, MAE, PSNR


if __name__ == '__main__':
    calculateSSIM()